import numpy as np 
import os
import skimage.io as io
import skimage.transform as trans
import numpy as np
from keras.models import *
from keras.layers import *
from keras.optimizers import *
from keras.callbacks import ModelCheckpoint, LearningRateScheduler
from keras import backend as keras
from keras.regularizers import l2
from keras.layers.normalization import BatchNormalization

def unet(pretrained_weights = None,input_size = (256,256,1)):
    inputs = Input(input_size)
    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)
    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)
    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)
    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)
    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)
    drop4 = Dropout(rate=0.5)(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)

    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)
    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)
    drop5 = Dropout(rate=0.5)(conv5)

    up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))
    merge6 = concatenate([drop4,up6], axis = 3)
    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)

    up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))
    merge7 = concatenate([conv3,up7], axis = 3)
    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)

    up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))
    merge8 = concatenate([conv2,up8], axis = 3)
    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)
    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)

    up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))
    merge9 = concatenate([conv1,up9], axis = 3)
    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)
    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
    conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
    conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)

    model = Model([inputs], [conv10])

    model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])
    
    if(pretrained_weights):
        model.load_weights(pretrained_weights)

    return model



def alexnet_model(img_shape=(256, 256), n_classes=5, l2_reg=0.,weights=None):

    # Initialize model
    alexnet = Sequential()
    ipt = Input([256,256,1])
    l1_conv2D = Conv2D(96, (11, 11), input_shape=img_shape,padding='same', kernel_regularizer=l2(l2_reg))(ipt)
    l1_bn = BatchNormalization()(l1_conv2D)
    l1_act = Activation('relu')(l1_bn)
    l1_pool = MaxPooling2D(pool_size=(2, 2))(l1_act)

    # Layer 2
    l2_conv2D = Conv2D(256, (5, 5), padding='same')(l1_pool)
    l2_bn = BatchNormalization()(l2_conv2D)
    l2_act = Activation('relu')(l2_bn)
    l2_pool = MaxPooling2D(pool_size=(2, 2))(l2_act)

    # Layer 3
    l3_pad = ZeroPadding2D((1, 1))(l2_pool)
    l3_conv2D = Conv2D(512, (3, 3), padding='same')(l3_pad)
    l3_bn = BatchNormalization()(l3_conv2D)
    l3_act = Activation('relu')(l3_bn)
    l3_pool = MaxPooling2D(pool_size=(2, 2))(l3_act)

    # Layer 4
    l4_pad = ZeroPadding2D((1, 1))(l3_pool)
    l4_conv2D = Conv2D(1024, (3, 3), padding='same')(l4_pad)
    l4_bn = BatchNormalization()(l4_conv2D)
    l4_act = Activation('relu')(l4_bn)

    # Layer 5
    l5_pad = ZeroPadding2D((1, 1))(l4_act)
    l5_conv2D = Conv2D(1024, (3, 3), padding='same')(l5_pad)
    l5_bn = BatchNormalization()(l5_conv2D)
    l5_act = Activation('relu')(l5_bn)
    l5_pool = MaxPooling2D(pool_size=(2, 2))(l5_act)

    # Layer 6
    l6_flt = Flatten()(l5_pool)
    l6_ds = Dense(3072)(l6_flt)
    l6_bn = BatchNormalization()(l6_ds)
    l6_act = Activation('relu')(l6_bn)
    l6_do = Dropout(rate=0.5)(l6_act)

    # Layer 7
    l7_ds = Dense(4096)(l6_do)
    l7_bn = BatchNormalization()(l7_ds)
    l7_act = Activation('relu')(l7_bn)
    l7_do = Dropout(rate=0.5)(l7_act)

    # Layer 8
    l8_ds = Dense(n_classes)(l7_do)
    l8_bn = BatchNormalization()(l8_ds)
    l8_act = Activation('softmax')(l8_bn)

    model = Model(ipt,l8_act)
    if weights is not None:
        model.load_weights(weights)

    return model

if __name__ == "__main__":
    alexnet_model()